Eager in/clone_in - #9277
Conversation
The old two-phase wrapper model recorded wrappers in a map and applied them during lowering (WrapCalls). Because resolution read the pre-rewrite graph while the rewrite was deferred, the literal and effective graphs diverged, which was the common cause of several bugs: - Issue 3661: cloning the same Func twice crashed, because deep-copying a Func that already carried wrappers couldn't remap them. - "Deletion via cloning": a clone_in that redirected the only path to an already-wrapped Func orphaned it, leaving it in the environment but dead in the effective graph, tripping an assert in RealizationOrder. Custom in(g)/clone_in now rewrite the named consumers eagerly, so the graph always reflects reality and both bugs become unreachable. The consumer is frozen afterwards, since a later definition wouldn't be wrapped. Global f.in() is expressed as a global_wrapper link on the Func plus a follow flag on call-node FunctionPtrs: get() follows the link, so a call resolves to the wrapper as if every caller had been rewritten, and the deep_copy that lowering already does materializes it (rebuilding each call with the wrapper's name). Self-references and wrapper bodies are marked not to follow (via WeakenFunctionPtrs) so they don't cycle. This is retroactive, future-capturing, and chains (f.in().in()) for free, and removes the need for the deferred WrapCalls pass, which is deleted. Adds a fuzz test combining in/clone_in into deep chains and indirect wraps, an error test for adding a definition after wrapping, and a func_clone regression test. Reorders two update-after-wrap tests to define-then-wrap. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
# Conflicts: # src/Func.cpp # src/IR.cpp # test/error/CMakeLists.txt # test/fuzz/CMakeLists.txt
Global .in() wrappers no longer occupy a "" entry in the wrappers map.
Idempotency is decided by Function::global_wrapper(), and the freeze +
follow-flag clearing that used to happen in add_wrapper("", W) now lives
in set_global_wrapper.
Function::global_wrapper() returns a strong, non-following handle to the
immediate wrapper. Following there would make f.in() resolve to the end
of the wrapper chain, so f.in().in() would wrap the wrong Func and, via
copy_to_host, hit "Extern Func has itself as an argument".
Serialize FunctionContents::global_wrapper as a WrapperRef so a
round-tripped pipeline keeps its global-wrapper links.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Describe the observable effect -- all past and future consumers are rewritten to call the wrapper -- instead of the follow-link mechanism. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
|
Looks like this breaks Anderson2021. |
|
Yeah, something aint right. Investigating. |
|
Autodiff was generating call nodes directly instead of using Func's operator(), and it was doing it without the new flag that says to follow global wrappers, so it wasn't seeing a global wrapper. A Func was then compute_at its own global wrapper, but that global wrapper had no consumers, so it wasn't in the pipeline, so that Func's schedule was invalid (compute_at something not in the pipeline). Fixing. |
The follow_global_wrappers flag is what makes a call node resolve to a Func's global wrapper (Func::in()). It was set only at the user-facing FuncRef chokepoints, so two consumer edges escaped it: - Autodiff builds its adjoint calls directly via Call::make, so a global wrapper never redirected them. The wrapper ended up with no consumers and was pruned, making a compute_at onto it an invalid location (seen in the anderson2021 cost-model schedule). - Pipeline::get_func returned a handle built from a Call node, inheriting its follow flag. After a first Func::in() the handle shifted to the new wrapper, so a second Func::in() wrapped the wrapper. Drop the default on Call::make's follow_global_wrappers argument so every consumer-edge site states intent, set it true in autodiff, and have get_func return a non-following handle. The rfactor self-reference and the ScheduleFunctions blend self-call stay non-following. Add tests: an rfactor+in test whose call graph pins the external edge (follows) versus the intermediate's self-reference (does not), and a test that Func::in() on a get_func handle is idempotent. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Call::make's Function overload no longer defaults follow_global_wrappers, so the printer must supply it. Emit op->func.follow_global_wrappers so the reconstructed call matches the signature and round-trips faithfully. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #9277 +/- ##
=======================================
Coverage 70.14% 70.14%
=======================================
Files 257 257
Lines 79162 79142 -20
Branches 18979 18978 -1
=======================================
- Hits 55529 55516 -13
+ Misses 17881 17806 -75
- Partials 5752 5820 +68 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
This changes clone_in and in to just eagerly rewrite all consumers to call via a wrapper (or copy, for clone_in), instead of maintaining a map of wrappers that needs to be reasoned about on all future calls to DAG-rewriting scheduling directives (e.g. eager_inline).
The argless Func::in() rewrites all existing consumers via a sleight of hand. Existing consumers hold a reference to this Function object via a FunctionPtr. This PR makes it so that when FunctionPtrs dereference to a Function (FunctionPtr::get()), they follow global wrapper chains instead of returning the first Function found, but only if a flag on the FunctionPtr is set. This means you can change the meaning of all existing FunctionPtrs that are defined with the flag set without a Function needing to keep a list of all its callers (which turns into a memory management nightmare - what if there's a caller in another thread that's currently going out of scope?)
Func::in with args doesn't have this problem, because the consumers to rewrite are handed to it as args, so they can just be rewritten then and there.
Also adds a fuzzer to try out weird in/clone_in combos to look for oddities. I think we can probably make this a more general Func-DAG-rewriting fuzzer in future.
Fixes #3661
Breaking changes
Changes the behavior of in and clone_in to be eager, but I can't think of any existing code that this would break in practice. People treat it that way already. It's observible to the user if they introspect on the RHS of a calling Func - they'll see that the rewrite has occurred already, instead of being done later during lowering.